import os
import h5py
import numpy as np
from scipy.ndimage import map_coordinates
from multiprocessing import Pool
from collections import deque
from easy_io import read_pkl_file, confirm_file_overwrite


def read_and_interp(scanid):
    with h5py.File('/ssd2/data/luna16_vol.h5', 'r') as f:
        vol = f[scanid][...]
    ori_spacing = spacing_dict[scanid]
    if ori_spacing[0] != ori_spacing[1]:
        print('\nSpacing of x and y mismatch: ', scanid)
    new_spacing = (min(ori_spacing),) * 3
    ori_shape = vol.shape
    vol = map_coordinates(
        vol,
        np.mgrid[
            0:ori_shape[0]:(new_spacing[0] / ori_spacing[0]),
            0:ori_shape[1]:(new_spacing[1] / ori_spacing[1]),
            0:ori_shape[2]:(new_spacing[2] / ori_spacing[2]),
        ],
        output='float32',
        order=1,
        mode='nearest',
    )
    np.savez_compressed(os.path.join('/ssd2/data/luna16_vol_interp', scanid), vol=vol)


def prepare_spacing_dict():
    global spacing_dict
    spacing_dict = read_pkl_file('/ssd2/data/luna16_origin_spacing.pkl')['spacing']


def gen():
    pool = Pool(processes=os.cpu_count() - 1, initializer=prepare_spacing_dict)
    queue = deque()
    with h5py.File('/ssd2/data/luna16_vol.h5', 'r') as f:
        scanlst = sorted(f.keys())
    for scanid in scanlst:
        queue.append(
            pool.apply_async(read_and_interp, args=(scanid,))
        )
    while len(queue) > 0:
        result = queue.popleft()
        result.get()


gen()

# def H5Writer(h5_file_path, mode, dataset_generator, attrs_dict=None):
#     if mode == 'a' or (mode == 'w' and confirm_file_overwrite(h5_file_path)):
#         g = h5py.File(h5_file_path, mode=mode, libver='latest')
#         try:
#             if attrs_dict:
#                 for k in attrs_dict:
#                     g.attrs[k] = attrs_dict[k]
#             for name, data in dataset_generator:
#                 g.create_dataset(name=name, data=data, compression='gzip')
#                 g.flush()
#         except:
#             g.close()
#             raise
#         else:
#             g.close()
#
# H5Writer('/ssd2/data/luna16_vol_interped.h5', 'w', gen())
